{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `N`amed `E`ntity `R`ecognition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from warnings import filterwarnings\n",
    "filterwarnings('ignore')  # 不打印警告"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>388568</th>\n",
       "      <td>NaN</td>\n",
       "      <td>impact</td>\n",
       "      <td>NN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388569</th>\n",
       "      <td>NaN</td>\n",
       "      <td>.</td>\n",
       "      <td>.</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388570</th>\n",
       "      <td>Sentence: 47959</td>\n",
       "      <td>Indian</td>\n",
       "      <td>JJ</td>\n",
       "      <td>B-gpe</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388571</th>\n",
       "      <td>NaN</td>\n",
       "      <td>forces</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388572</th>\n",
       "      <td>NaN</td>\n",
       "      <td>said</td>\n",
       "      <td>VBD</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388573</th>\n",
       "      <td>NaN</td>\n",
       "      <td>they</td>\n",
       "      <td>PRP</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388574</th>\n",
       "      <td>NaN</td>\n",
       "      <td>responded</td>\n",
       "      <td>VBD</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388575</th>\n",
       "      <td>NaN</td>\n",
       "      <td>to</td>\n",
       "      <td>TO</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388576</th>\n",
       "      <td>NaN</td>\n",
       "      <td>the</td>\n",
       "      <td>DT</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>388577</th>\n",
       "      <td>NaN</td>\n",
       "      <td>attack</td>\n",
       "      <td>NN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             Sentence #       Word  POS    Tag\n",
       "388568              NaN     impact   NN      O\n",
       "388569              NaN          .    .      O\n",
       "388570  Sentence: 47959     Indian   JJ  B-gpe\n",
       "388571              NaN     forces  NNS      O\n",
       "388572              NaN       said  VBD      O\n",
       "388573              NaN       they  PRP      O\n",
       "388574              NaN  responded  VBD      O\n",
       "388575              NaN         to   TO      O\n",
       "388576              NaN        the   DT      O\n",
       "388577              NaN     attack   NN      O"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = pd.read_csv('train.csv')\n",
    "df_test = pd.read_csv('test.csv')\n",
    "df_test.tail(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>659992</th>\n",
       "      <td>Sentence: 30159</td>\n",
       "      <td>control</td>\n",
       "      <td>VB</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>659993</th>\n",
       "      <td>Sentence: 30159</td>\n",
       "      <td>the</td>\n",
       "      <td>DT</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>659994</th>\n",
       "      <td>Sentence: 30159</td>\n",
       "      <td>island</td>\n",
       "      <td>NN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>659995</th>\n",
       "      <td>Sentence: 30159</td>\n",
       "      <td>chain</td>\n",
       "      <td>NN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>659996</th>\n",
       "      <td>Sentence: 30159</td>\n",
       "      <td>.</td>\n",
       "      <td>.</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             Sentence #     Word POS Tag\n",
       "659992  Sentence: 30159  control  VB   O\n",
       "659993  Sentence: 30159      the  DT   O\n",
       "659994  Sentence: 30159   island  NN   O\n",
       "659995  Sentence: 30159    chain  NN   O\n",
       "659996  Sentence: 30159        .   .   O"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = df_train.fillna(method='ffill')\n",
    "df_test = df_test.fillna(method='ffill')\n",
    "df_train.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Tag</th>\n",
       "      <th>count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>B-art</td>\n",
       "      <td>260</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B-eve</td>\n",
       "      <td>218</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>B-geo</td>\n",
       "      <td>23605</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>B-gpe</td>\n",
       "      <td>10045</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>B-nat</td>\n",
       "      <td>121</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>B-org</td>\n",
       "      <td>12507</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>B-per</td>\n",
       "      <td>10681</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>B-tim</td>\n",
       "      <td>12756</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>I-art</td>\n",
       "      <td>197</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>I-eve</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>I-geo</td>\n",
       "      <td>4707</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>I-gpe</td>\n",
       "      <td>133</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>I-nat</td>\n",
       "      <td>35</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>I-org</td>\n",
       "      <td>10384</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>I-per</td>\n",
       "      <td>11000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>I-tim</td>\n",
       "      <td>4125</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>O</td>\n",
       "      <td>559046</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      Tag   count\n",
       "0   B-art     260\n",
       "1   B-eve     218\n",
       "2   B-geo   23605\n",
       "3   B-gpe   10045\n",
       "4   B-nat     121\n",
       "5   B-org   12507\n",
       "6   B-per   10681\n",
       "7   B-tim   12756\n",
       "8   I-art     197\n",
       "9   I-eve     177\n",
       "10  I-geo    4707\n",
       "11  I-gpe     133\n",
       "12  I-nat      35\n",
       "13  I-org   10384\n",
       "14  I-per   11000\n",
       "15  I-tim    4125\n",
       "16      O  559046"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = df_train.Tag.unique().tolist()\n",
    "labels.remove('O')\n",
    "df_train.groupby('Tag').size().reset_index(name='count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(659997,) (388578,)\n"
     ]
    }
   ],
   "source": [
    "y_train, y_test = df_train.Tag.values, df_test.Tag.values\n",
    "print(y_train.shape, y_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**多数票决**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = df_train.Word.tolist()\n",
    "X_test = df_test.Word.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Majority_vote: \n",
    "\n",
    "    def fit(self, X, y):\n",
    "        counter = {}\n",
    "        for w, t in zip(X, y):\n",
    "            if w in counter:\n",
    "                if t in counter[w]:\n",
    "                    counter[w][t] += 1\n",
    "                else:\n",
    "                    counter[w][t] = 1\n",
    "            else:\n",
    "                counter[w] = {t: 1}\n",
    "        self.vote = {}\n",
    "        for w, t in counter.items():\n",
    "            self.vote[w] = max(t, key=t.get)\n",
    "        return self\n",
    "\n",
    "    def predict(self, X):\n",
    "        return [self.vote.get(x, 'O') for x in X]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "clf = Majority_vote().fit(X_train, y_train)\n",
    "y_pred = clf.predict(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             precision    recall  f1-score   support\n",
      "\n",
      "      B-geo       0.79      0.85      0.82     14039\n",
      "      B-gpe       0.94      0.95      0.95      5825\n",
      "      B-per       0.79      0.63      0.70      6309\n",
      "      I-geo       0.70      0.61      0.65      2707\n",
      "      B-org       0.68      0.48      0.57      7636\n",
      "      I-org       0.71      0.52      0.60      6400\n",
      "      B-tim       0.87      0.76      0.81      7577\n",
      "      B-art       0.19      0.06      0.09       142\n",
      "      I-art       0.03      0.01      0.01       100\n",
      "      I-per       0.72      0.65      0.68      6251\n",
      "      I-gpe       0.53      0.65      0.58        65\n",
      "      I-tim       0.61      0.12      0.20      2403\n",
      "      B-nat       0.45      0.46      0.45        80\n",
      "      B-eve       0.61      0.22      0.33        90\n",
      "      I-eve       0.33      0.09      0.14        76\n",
      "      I-nat       0.00      0.00      0.00        16\n",
      "\n",
      "avg / total       0.77      0.68      0.71     59716\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "report = classification_report(y_test, y_pred, labels)\n",
    "print(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**条件随机场**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def groupby(s):\n",
    "    f = lambda s: [(w, p, t) for w, p, t in zip(\n",
    "        s.Word.values, s.POS.values, s.Tag.values)]\n",
    "    return list(s.groupby('Sentence #').apply(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('Thousands', 'NNS', 'O'),\n",
       "  ('of', 'IN', 'O'),\n",
       "  ('demonstrators', 'NNS', 'O'),\n",
       "  ('have', 'VBP', 'O'),\n",
       "  ('marched', 'VBN', 'O'),\n",
       "  ('through', 'IN', 'O'),\n",
       "  ('London', 'NNP', 'B-geo'),\n",
       "  ('to', 'TO', 'O'),\n",
       "  ('protest', 'VB', 'O'),\n",
       "  ('the', 'DT', 'O'),\n",
       "  ('war', 'NN', 'O'),\n",
       "  ('in', 'IN', 'O'),\n",
       "  ('Iraq', 'NNP', 'B-geo'),\n",
       "  ('and', 'CC', 'O'),\n",
       "  ('demand', 'VB', 'O'),\n",
       "  ('the', 'DT', 'O'),\n",
       "  ('withdrawal', 'NN', 'O'),\n",
       "  ('of', 'IN', 'O'),\n",
       "  ('British', 'JJ', 'B-gpe'),\n",
       "  ('troops', 'NNS', 'O'),\n",
       "  ('from', 'IN', 'O'),\n",
       "  ('that', 'DT', 'O'),\n",
       "  ('country', 'NN', 'O'),\n",
       "  ('.', '.', 'O')],\n",
       " [('Iranian', 'JJ', 'B-gpe'),\n",
       "  ('officials', 'NNS', 'O'),\n",
       "  ('say', 'VBP', 'O'),\n",
       "  ('they', 'PRP', 'O'),\n",
       "  ('expect', 'VBP', 'O'),\n",
       "  ('to', 'TO', 'O'),\n",
       "  ('get', 'VB', 'O'),\n",
       "  ('access', 'NN', 'O'),\n",
       "  ('to', 'TO', 'O'),\n",
       "  ('sealed', 'JJ', 'O'),\n",
       "  ('sensitive', 'JJ', 'O'),\n",
       "  ('parts', 'NNS', 'O'),\n",
       "  ('of', 'IN', 'O'),\n",
       "  ('the', 'DT', 'O'),\n",
       "  ('plant', 'NN', 'O'),\n",
       "  ('Wednesday', 'NNP', 'B-tim'),\n",
       "  (',', ',', 'O'),\n",
       "  ('after', 'IN', 'O'),\n",
       "  ('an', 'DT', 'O'),\n",
       "  ('IAEA', 'NNP', 'B-org'),\n",
       "  ('surveillance', 'NN', 'O'),\n",
       "  ('system', 'NN', 'O'),\n",
       "  ('begins', 'VBZ', 'O'),\n",
       "  ('functioning', 'VBG', 'O'),\n",
       "  ('.', '.', 'O')]]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = groupby(df_train)\n",
    "df_test = groupby(df_test)\n",
    "df_train[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Processing:\n",
    "\n",
    "    def __init__(self, tpl):\n",
    "        self.tpl = tpl\n",
    "        self.len = len(tpl)\n",
    "\n",
    "    def get_features(self, i):\n",
    "        word = self.tpl[i][0]\n",
    "        postag = self.tpl[i][1]\n",
    "\n",
    "        features = {\n",
    "            'bias': 1.0,\n",
    "            'word.lower()': word.lower(),\n",
    "            'word[-3:]': word[-3:],\n",
    "            'word[-2:]': word[-2:],\n",
    "            'word.isupper()': word.isupper(),\n",
    "            'word.istitle()': word.istitle(),\n",
    "            'word.isdigit()': word.isdigit(),\n",
    "            'postag': postag,\n",
    "            'postag[:2]': postag[:2],\n",
    "        }\n",
    "        if i > 0:\n",
    "            word1 = self.tpl[i - 1][0]\n",
    "            postag1 = self.tpl[i - 1][1]\n",
    "            features.update({\n",
    "                '-1:word.lower()': word1.lower(),\n",
    "                '-1:word.istitle()': word1.istitle(),\n",
    "                '-1:word.isupper()': word1.isupper(),\n",
    "                '-1:postag': postag1,\n",
    "                '-1:postag[:2]': postag1[:2],\n",
    "            })\n",
    "        else:\n",
    "            features['BOS'] = True\n",
    "        if i < self.len - 1:\n",
    "            word1 = self.tpl[i + 1][0]\n",
    "            postag1 = self.tpl[i + 1][1]\n",
    "            features.update({\n",
    "                '+1:word.lower()': word1.lower(),\n",
    "                '+1:word.istitle()': word1.istitle(),\n",
    "                '+1:word.isupper()': word1.isupper(),\n",
    "                '+1:postag': postag1,\n",
    "                '+1:postag[:2]': postag1[:2],\n",
    "            })\n",
    "        else:\n",
    "            features['EOS'] = True\n",
    "\n",
    "        return features\n",
    "\n",
    "    def to_labels(self):\n",
    "        return [l for t, p, l in self.tpl]\n",
    "\n",
    "    def to_features(self):\n",
    "        return [self.get_features(i) for i in range(self.len)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, y_train = [], []\n",
    "for tpl in df_train:\n",
    "    tpl = Processing(tpl)\n",
    "    X_train.append(tpl.to_features())\n",
    "    y_train.append(tpl.to_labels())\n",
    "X_test, y_test = [], []\n",
    "for tpl in df_train:\n",
    "    tpl = Processing(tpl)\n",
    "    X_test.append(tpl.to_features())\n",
    "    y_test.append(tpl.to_labels())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn_crfsuite import CRF\n",
    "clf = CRF(\n",
    "    algorithm='lbfgs',\n",
    "    c1=.1,\n",
    "    c2=.1,\n",
    "    max_iterations=100,\n",
    ").fit(X_train, y_train)\n",
    "y_pred = clf.predict(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             precision    recall  f1-score   support\n",
      "\n",
      "      B-geo       0.91      0.96      0.93     23605\n",
      "      B-gpe       0.98      0.95      0.96     10045\n",
      "      B-per       0.96      0.93      0.94     10681\n",
      "      I-geo       0.91      0.93      0.92      4707\n",
      "      B-org       0.92      0.87      0.89     12507\n",
      "      I-org       0.95      0.94      0.94     10384\n",
      "      B-tim       0.97      0.93      0.95     12756\n",
      "      B-art       0.94      0.77      0.85       260\n",
      "      I-art       0.92      0.81      0.86       197\n",
      "      I-per       0.95      0.96      0.95     11000\n",
      "      I-gpe       0.97      0.66      0.79       133\n",
      "      I-tim       0.94      0.90      0.92      4125\n",
      "      B-nat       0.84      0.64      0.72       121\n",
      "      B-eve       0.88      0.78      0.82       218\n",
      "      I-eve       0.89      0.74      0.81       177\n",
      "      I-nat       0.89      0.71      0.79        35\n",
      "\n",
      "avg / total       0.94      0.93      0.94    100951\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn_crfsuite.metrics import flat_classification_report\n",
    "report = flat_classification_report(y_test, y_pred, labels)\n",
    "print(report)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
